"""
Created on 10-Dec-2009

@author: dexen
"""



import os
import sys
import time
import multiprocessing as mp


from dexen.system.base.spec import DexenSpec
from dexen.system.server.proxy import SlaveProxy
from dexen.system.node.master.taskq_mgr import TaskQMgr
from dexen.system.server.info import ExecInfo, SlaveInfo
from dexen.system.base.task import OutputUnit, InputUnit, TaskInfo
from dexen.system import util
import rpyc
    
#===============================================================================
# 
#===============================================================================
class SlaveMgr(mp.Process):
    def __init__(self, spec, comm, reg_addr):
        mp.Process.__init__(self)
        self.spec = spec
        self.comm = comm
        self.reg_addr = reg_addr
        self.slaves = {} # proxies
        assert isinstance(self.spec, DexenSpec)
        self.tq_mgr = TaskQMgr(spec.tasks)
        self.cur_unit = None
        self.sent_tasks = {}
         
    def _setup_schema(self):
        sys.path.append(self.spec.masterdir)
        util.unzip(self.spec.schema_zip, self.spec.masterdir) # TODO: consider folder structure
        
    def _init_job(self):
        """
        Executes the init task and puts the output into the queues.
        Thus, this method kickstarts the job execution
        """
        print "scheduler->init(), ", self.spec.tasks
        masterdir = self.spec.masterdir
        task = self.spec.init_task
        assert isinstance(task, TaskInfo)
        # assert task is init type TODO:
        prevpath = os.getcwd()
        os.chdir(masterdir)                    
        obj = task.get_obj()
        inds = obj.__initialize__()
        self.tq_mgr.set_inds(task.name, inds)
        os.chdir(prevpath)
        print "the end of init(spec), ", self.tq_mgr
    
    def _recover(self, slave):
        """
        This function is called from 3 places.
        the first connection (during proxy object creation), execute, schedule
        """
        input = slave.input
        """ 
        input is None only when the crash occurs during proxy creation
        """
        if input != None: 
            assert isinstance(input, InputUnit)
            self.tq_mgr.recover_inds(input.task_name, input.inds)
        slave_info = slave.get_info()
        self.comm.update_slave(slave_info)
                       
    def _get_slave_addr(self, slave_name):
        conn = rpyc.connect(self.reg_addr[0], self.reg_addr[1])
        addr = conn.root.get_slave_addr(slave_name)
        conn.close()
        return addr
    
    def _schedule(self, slave_name, input_unit):
        assert isinstance(input_unit, InputUnit)
        if not self.slaves.has_key(slave_name): 
            addr = self._get_slave_addr(slave_name)
            slave = SlaveProxy(slave_name, self.spec, addr)
            if slave.is_crashed():
                self._recover(slave)
            self.slaves[slave_name] = slave
        
        slave = self.slaves[slave_name]
        assert isinstance(slave, SlaveProxy)
        
        print "executing", slave.name, time.time()
        slave.execute(input_unit)
        if slave.is_crashed():
            self._recover(slave)

    # TODO: cater for crashes
    def _update_slaves(self):
        #print "updating slaves"
        for slave in self.slaves.values():
            assert isinstance(slave, SlaveProxy)
            
            if not slave.is_executing():
                continue

            slave.update()
    
            if slave.is_executing():
                # This means that after querying the slave, the slave hasn't 
                # finished executing the task.
                continue
            
            if slave.is_crashed():
                """ Note that the slave.inds is tuple (task_name, inds)
                This is not the crash of task, this is the crash of slave
                The info about task crash is kept in the outputunit 
                """
                print slave.name, "crashed"
                self._recover(slave)

            if slave.is_idle():
                """
                If the slave is in idle state, it means that the proxy has just
                got back the results pending to be processed
                """
                o_unit = slave.output
                assert isinstance(o_unit, OutputUnit)
                """
                TODO: you may want to analyze the contents of o_unit 
                (e.g. crashed or not execution time, vs..)
                """

                print slave.name, "has become idle", time.time()

                exec_info = ExecInfo(o_unit.task_name)
                
                if o_unit.is_crashed():
                    """
                    TODO: you might wanna recover the inds
                    """
                    exec_info.crashed()
                    print "execution is crashed in", slave.name, time.time()
                else:  
                    self.tq_mgr.set_inds(o_unit.task_name, o_unit.inds)
                    exec_info.set_exec_time(o_unit.exec_time)
                    
                slave_info = slave.get_info()
                assert isinstance(slave_info, SlaveInfo)
                slave_info.exec_info = exec_info
                self.comm.update_slave(slave_info)
    
    def run(self):
        self._setup_schema()
        self._init_job()
        
        while True:
            self._update_slaves()
            
            if not self.comm.has_unit():
                task_name, inds = self.tq_mgr.get_inds()
                if not task_name: 
                    continue
                self.comm.has_unit(True)
                input_unit = InputUnit(task_name, inds)
                
            slave_name = self.comm.get_slave()
            if not slave_name: 
                continue
            self._schedule(slave_name, input_unit)    


